#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Copyright 2020 Huawei Technologies Co., Ltd
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
import acl
import sys
import copy
from .thread_pool import ThreadPool
from .mem_pool import SimpleMemoryPool

import os
import threading

sys.path.append("..")
from common.const import DEBUG, INFO, WARNING, ERROR, SET_DEVICE, RESET_DEVICE, RUN_MODE
from common.log import Log, check_ret

"""
EXAMPLE:
ResourceMangerInfo:

device_info = [0,1,2,3]

device_context_info = {0:[contex0], 
                       1:[contex1],
                       2:[contex2],
                       3:[contex3]}

context_stream_info = {contex0:[stream00,stream01],
                       contex1:[stream10,stream11],
                       contex2:[stream20,stream21],
                       contex3:[stream30,stream31]}
"""


class ResourceManagerInfo():
    def __init__(self):
        self.device_info = []
        self.device_context_info = {}
        self.context_stream_info = {}
        self.locker = threading.Lock()

    def add_device_info(self, device_id):
        self.device_info.append(device_id)

    def add_device_context_info(self, device_id, context_id):
        self.locker.acquire()
        if device_id in self.device_context_info.keys():
            if context_id not in self.device_context_info[device_id]:
                self.device_context_info[device_id].append(context_id)
            else:
                print("Exception: why has context_id already existed in the list ?")
        else:
            self.device_context_info.update({device_id: [context_id]})
        self.locker.release()

    def add_context_stream_info(self, context_id, stream_id):
        self.locker.acquire()
        if context_id in self.context_stream_info.keys():
            if stream_id not in self.context_stream_info[context_id]:
                self.context_stream_info[context_id].append(stream_id)
            else:
                print("Exception: why has stream_id already existed in the list ?")
        else:
            self.context_stream_info.update({context_id: [stream_id]})
        self.locker.release()

    def del_device_info(self, device_id):
        self.locker.acquire()
        if device_id in self.device_info:
            self.device_info.remove(device_id)
        else:
            print('Exception: no device_id in device_info ?')
        self.locker.release()

    def del_context_info(self, context_id):
        self.locker.acquire()
        for device_id in self.device_info:
            if context_id in self.device_context_info[device_id]:
                self.device_context_info[device_id].remove(context_id)
                break
        self.locker.release() 

    def del_stream_info(self, stream_id):
        self.locker.acquire()
        for context_id in self.context_stream_info.keys():
            if stream_id in self.context_stream_info[context_id]:
                self.context_stream_info[context_id].remove(stream_id)
        self.locker.release()



class AclRuntimeManger():
    def __init__(self):
        self.class_name = self.__class__.__name__
        self.device_count = 0
        self.device_on_line = self.device_check()
        self.run_mode = None

    def init_acl_resource(self, path=None):
        if path is None:
            ret = acl.init()
            check_ret(self.class_name, 'init_acl_resource', 'fail to init acl', ret)
        else:
            if os.path.exists(path):
                ret = acl.init(path)
                check_ret(self.class_name, 'init_acl_resource', 'fail to init acl', ret)
            else:
                check_ret(self.class_name, 'init_acl_resource', \
                    'fail to open path', 1)

    def device_check(self):
        _, ret = acl.rt.get_device_count()
        check_ret(self.class_name, 'device_check', 'fail to get device', ret)
        return

    def device_runtime(self, device_id:int, flag:int):
        if flag == SET_DEVICE:
            ret = acl.rt.set_device(device_id)
            check_ret(self.class_name, 'device_runtime', 'fail to set device_id', ret)
        elif flag == RESET_DEVICE:
            ret = acl.rt.reset_device(device_id)
            check_ret(self.class_name, 'device_runtime', 'fail to reset device_id', ret)
        elif flag == RUN_MODE:
            run_mode, ret = acl.rt.get_run_mode()
            check_ret(self.class_name, 'device_runtime', 'fail to get run_mode', ret)
            return run_mode
        else:
            check_ret(self.class_name, 'device_runtime', 'flag value is out of its range', 1)

    def finalize_acl_resource(self):
        ret = acl.finalize()
        check_ret(self.class_name, 'finalize_acl_resource', 'fail to finalize', ret)


class ResourceManager():
    def __init__(self):
        self.class_name = self.__class__.__name__
        self.runtime_manage = AclRuntimeManger()
        self.info = ResourceManagerInfo()

    def init_resource_manager(self, acl_json_config_path=None):
        self.runtime_manage.init_acl_resource(path = acl_json_config_path)
            
    def device_combination(self, device_id):
        self.runtime_manage.device_runtime(device_id, SET_DEVICE)
        self.info.add_device_info(device_id)

    def create_context(self, device_id):
        context_id, ret = acl.rt.create_context(device_id)
        check_ret(self.class_name, 'create_context', 'fail to create context', ret)
        self.info.add_device_context_info(device_id, context_id)
        return context_id, ret

    def create_stream(self, context_id):
        ret = acl.rt.set_context(context_id)
        stream_id, ret = acl.rt.create_stream()
        check_ret(self.class_name, 'create_stream', 'fail to create stream', ret)
        self.info.add_context_stream_info(context_id, stream_id)
        return stream_id, ret

    def set_current_context(self, context_id):
        ret = acl.rt.set_context(context_id)
        check_ret(self.class_name, 'set_current_context', 'fail to set context', ret)

    def destroy_stream(self, stream_id):
        ret = acl.rt.destroy_stream(stream_id)
        check_ret(self.class_name, 'destroy_stream', 'fail to destroy stream', ret)
        self.info.del_stream_info(stream_id)

    def destroy_context(self, context_id):
        ret = acl.rt.destroy_context(context_id)
        check_ret(self.class_name, 'destroy_context', 'fail to destroy context', ret)
        self.info.del_context_info(context_id)

    def reset_device(self, device_id):
        ret = acl.rt.reset_device(device_id)
        check_ret(self.class_name, 'reset_device', 'fail to reset device', ret)

    def release_device_resource(self,  device_id):
        self.runtime_manage.device_runtime(device_id, RESET_DEVICE)
        self.info.del_device_info(device_id)

    def release_runtime_resource(self):
        device_info = copy.deepcopy(self.info.device_info)
        device_context_info  = copy.deepcopy(self.info.device_context_info)
        context_stream_info = copy.deepcopy(self.info.context_stream_info)
      
        if len(device_info) == 0:
            return
        for device_id in device_info:
            if device_id in device_context_info.keys():
                if len(device_context_info[device_id]) == 0:
                    Log.do_log(INFO, self.class_name, 'release_runtime_resource', \
                            'no used context in device[{}]'.format(device_id))
                else:
                    for context_id in device_context_info[device_id]:
                        if context_id in context_stream_info.keys():
                            self.set_current_context(context_id)
                            if len(context_stream_info[context_id]) == 0:
                                Log.do_log(INFO, self.class_name, 'release_runtime_resource', \
                                        'no used context in device[{}]'.format(device_id))
                            else:
                                for stream_id in context_stream_info[context_id]:
                                    self.destroy_stream(stream_id)
                        self.destroy_context(context_id)
                self.release_device_resource(device_id)
        self.runtime_manage.finalize_acl_resource()

    def stream_constraint_check(self):
        pass

    def context_contraint_check(self):
        pass

    def __del__(self):
        self.release_runtime_resource()


class AscendResource():
    def __init__(self):
        self.thread_pool = ThreadPool()
        self.resource_manager = ResourceManager()
        self.memory_pool = SimpleMemoryPool()
        self.resource_info = self.resource_manager.info
        self.host_thread_info = self.thread_pool.host_thread_pool_info

ascend_resource = AscendResource()

